/*!
  @file disparity_to_cloud.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include <vector>

#include <ros/ros.h>

#include <cv_bridge/cv_bridge.h>
#include <image_geometry/stereo_camera_model.h>
#include <image_transport/image_transport.h>
#include <image_transport/subscriber_filter.h>
#include <message_filters/subscriber.h> 
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <tf/transform_listener.h>

#include <sensor_msgs/Image.h>
#include <stereo_msgs/DisparityImage.h>
#include <sensor_msgs/PointCloud2.h>
#include <nav_msgs/OccupancyGrid.h>
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>

#include <Eigen/Dense>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/objdetect/objdetect.hpp>

#include <pcl/common/transforms.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/PointIndices.h>
#include <pcl/ros/conversions.h>
#include <pcl/octree/octree.h>

#include <nodelet/nodelet.h>

namespace stereo 
{

/*! 
 *  @brief  Converter from disparity image to point cloud
 *  @author Kyohei Otsu <kyon@ac.jaxa.jp>
 *  @date   2013
 *  
 */

class ConverterCloud : public nodelet::Nodelet
{
 public:
  typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image, stereo_msgs::DisparityImage> SyncPolicy;
  typedef pcl::PointXYZRGB PointType;
  typedef pcl::PointCloud<PointType> Cloud;
  typedef typename Cloud::Ptr CloudPtr;
  typedef pcl::octree::OctreePointCloudSearch<PointType> OctreeSearch;
  typedef std::vector<int> Indices;
  typedef typename boost::shared_ptr<Indices> IndicesPtr;


  virtual void onInit();

  void imageCB(const sensor_msgs::ImageConstPtr &imgL_msg, 
               const stereo_msgs::DisparityImageConstPtr &disp_msg);

 private:
  void cloud2map(const CloudPtr &cloud, nav_msgs::OccupancyGrid &map);

  void initTF();

  void recogFaces(const cv::Mat &img);

  template <typename T>
  float conv_bgr(const T &bgr)
  {
    uint8_t b = bgr[0];
    uint8_t g = bgr[1];
    uint8_t r = bgr[2];
    uint32_t rgb = ((uint32_t)r << 16 | (uint32_t)g << 8 | (uint32_t)b);
    return *reinterpret_cast<float*>(&rgb);
  }

  ros::NodeHandle nh_;
  ros::NodeHandle private_nh_;
  boost::shared_ptr<image_transport::ImageTransport> it_;
  image_transport::SubscriberFilter imgL_sub_;
  boost::shared_ptr<message_filters::Synchronizer<SyncPolicy> > sync_;
  message_filters::Subscriber<stereo_msgs::DisparityImage> disp_sub_;
  ros::Publisher cloud_pub_;
  ros::Publisher map_pub_;
  ros::Publisher marker_pub_;

  cv::Mat Q_;
  nav_msgs::OccupancyGrid map_;
  Eigen::Affine3f bTc_;
};

void ConverterCloud::onInit()
{
  nh_ = getNodeHandle();
  //it_ = image_transport::ImageTransport(nh_);
  it_.reset(new image_transport::ImageTransport(nh_));
  imgL_sub_ .subscribe(*it_, nh_.resolveName("left/image_rect_color"), 10);
  disp_sub_ .subscribe(nh_, nh_.resolveName("disparity/image"), 5);
  sync_.reset(new message_filters::Synchronizer<SyncPolicy>(SyncPolicy(10), imgL_sub_, disp_sub_));
  sync_->registerCallback(boost::bind(&ConverterCloud::imageCB, this, _1, _2));

  cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>(nh_.resolveName("disparity/cloud"), 3);
  map_pub_ = nh_.advertise<nav_msgs::OccupancyGrid>(nh_.resolveName("disparity/map"), 3);
  marker_pub_ = nh_.advertise<visualization_msgs::MarkerArray>(nh_.resolveName("disparity/dem"), 3);


  initTF();
}

void ConverterCloud::imageCB(
    const sensor_msgs::ImageConstPtr &imgL_msg, 
    const stereo_msgs::DisparityImageConstPtr &disp_msg)
{ 
  //if (cloud_pub_.getNumSubscribers() > 0 || map_pub_.getNumSubscribers() > 0)
  {
    cv::Mat imgL_orig = cv_bridge::toCvShare(imgL_msg)->image;
    cv::Mat data = cv_bridge::toCvShare(disp_msg->image, disp_msg)->image;
    cv::Mat imgL(data.rows, data.cols, CV_8UC3);
    cv::resize(imgL_orig, imgL, imgL.size());

    if (!Q_.data)
    {
      Q_ = cv::Mat::zeros(4, 4, CV_64F);
      Q_.at<double>(0, 0) = Q_.at<double>(1, 1) = 1.0;
      Q_.at<double>(0, 3) = - imgL.cols / 2;
      Q_.at<double>(1, 3) = - imgL.rows / 2;
      Q_.at<double>(2, 3) =   disp_msg->f;
      Q_.at<double>(3, 2) = 1.0 / disp_msg->T;
      ROS_INFO_STREAM("Q " << std::endl << Q_);
    }


    cv::Mat depth_cloud(data.rows, data.cols, CV_32FC3);
    cv::reprojectImageTo3D(data, depth_cloud, Q_, false, CV_32F);
    sensor_msgs::PointCloud2Ptr cloud_msg(new sensor_msgs::PointCloud2);
    CloudPtr cloud(new Cloud);
    //cloud->width = data.rows * data.cols;
    cloud->height = 1;
    cloud->is_dense = false;
    //cloud->points.resize(cloud->width * cloud->height);
    int step = data.cols;
    for (int i = 0; i < depth_cloud.rows; ++i)
      for (int j = 0; j < depth_cloud.cols; ++j)
      {
        cv::Vec3f &pt = depth_cloud.at<cv::Vec3f>(i, j);
        if (pt[2] != image_geometry::StereoCameraModel::MISSING_Z
            && !std::isinf(pt[2]) && pt[2] < 8)
        {
          PointType pt;
          pt.x = depth_cloud.at<cv::Vec3f>(i, j)[0];
          pt.y = depth_cloud.at<cv::Vec3f>(i, j)[1];
          pt.z = depth_cloud.at<cv::Vec3f>(i, j)[2];
          pt.rgb = conv_bgr(imgL.at<cv::Vec3b>(i, j));
          cloud->points.push_back(pt);
          /*
          cloud->points[i * step + j].x = depth_cloud.at<cv::Vec3f>(i, j)[0];
          cloud->points[i * step + j].y = depth_cloud.at<cv::Vec3f>(i, j)[1];
          cloud->points[i * step + j].z = depth_cloud.at<cv::Vec3f>(i, j)[2];
          cloud->points[i * step + j].rgb = conv_bgr(imgL.at<cv::Vec3b>(i, j));
          */
        }
        {
          //cloud->points[i * width + j].x = 1000;
          //cloud->points[i * width + j].y = 1000;
          //cloud->points[i * width + j].z = 1000;
        }
      }
    cloud->width = cloud->points.size();
    ROS_INFO_STREAM("published points: " << cloud->points.size());
    
    //if (cloud_pub_.getNumSubscribers() > 0)
    if (0)
    {
      pcl::toROSMsg(*cloud, *cloud_msg);
      cloud_msg->header = imgL_msg->header;
      cloud_msg->header.stamp = ros::Time::now();  // tmp: refleshing time stamp
      cloud_pub_.publish(cloud_msg);
    }


    //if (map_pub_.getNumSubscribers() > 0)
    {
      map_.header.stamp = ros::Time::now();
      map_.header.frame_id = "/map";
      cloud2map(cloud, map_);
      map_pub_.publish(map_);
    }


    /* face recognition */
    std::vector<cv::Rect_<int> > faces;
    cv::Mat gray(imgL_orig.rows, imgL_orig.cols, CV_8U);
    cv::Mat gray_big(768, 1024, CV_8U);
    cv::cvtColor(imgL_orig, gray, CV_BGR2GRAY);
    cv::resize(gray, gray_big, gray_big.size());
    //recogFaces(gray_big, faces);

    static cv::CascadeClassifier cascade_("/opt/ros/groovy/share/OpenCV/haarcascades/haarcascade_frontalface_default.xml");
    cascade_.detectMultiScale(gray_big, faces, 1.1, 3, 0, cv::Size(), cv::Size());
    {
      visualization_msgs::MarkerArrayPtr markers(new visualization_msgs::MarkerArray);
      visualization_msgs::Marker m;
      m.header.frame_id = "/map";
      m.header.stamp = ros::Time::now();
      m.type = visualization_msgs::Marker::CYLINDER;
      m.action = visualization_msgs::Marker::ADD;
      m.scale.x = 0.3;
      m.scale.y = 0.5;
      m.scale.z = 1.2;//0.01;
      m.color.a = 0.4; //1.0;
      m.color.r = 0.0;
      m.color.g = 0.0;
      m.color.b = 1.0;
      m.lifetime = ros::Duration(3);
      m.id = 0;

      float sc = 1.0 * data.rows / gray_big.rows;

      CloudPtr cloud_mark(new Cloud);
      CloudPtr cloud_mark_map(new Cloud);
      for (size_t i = 0; i < faces.size(); ++i)
      {
        ++m.id;
        PointType pt;
        pt.x = depth_cloud.at<cv::Vec3f>((faces[i].y+faces[i].height)*sc, (faces[i].x+0.5*faces[i].width)*sc)[0];
        pt.y = depth_cloud.at<cv::Vec3f>((faces[i].y+faces[i].height)*sc, (faces[i].x+0.5*faces[i].width)*sc)[1];
        pt.z = depth_cloud.at<cv::Vec3f>((faces[i].y+faces[i].height)*sc, (faces[i].x+0.5*faces[i].width)*sc)[2];
        cloud_mark->points.push_back(pt);
        /*
        m.pose.position.x = depth_cloud.at<cv::Vec3f>((faces[i].y+0.5*faces[i].height+100)*sc, (faces[i].x+0.5*faces[i].width)*sc)[0];
        m.pose.position.y = depth_cloud.at<cv::Vec3f>((faces[i].y+0.5*faces[i].height+100)*sc, (faces[i].x+0.5*faces[i].width)*sc)[1];
        m.pose.position.z = depth_cloud.at<cv::Vec3f>((faces[i].y+0.5*faces[i].height+100)*sc, (faces[i].x+0.5*faces[i].width)*sc)[2];
        */
        //m.pose.position.z = 0.1;
        //if (m.pose.position.x <= 0.1) continue;
        //if (m.pose.position.x == image_geometry::StereoCameraModel::MISSING_Z) continue;
        //markers->markers.push_back(m);

        //cv::rectangle(gray_big, faces[i], cv::Scalar(255, 255, 255));
      }
      pcl::transformPointCloud(*cloud_mark, *cloud_mark_map, bTc_);
      for (size_t i = 0; i < cloud_mark_map->points.size(); ++i)
      {
        m.pose.position.x = cloud_mark_map->points[i].x;
        m.pose.position.y = cloud_mark_map->points[i].y;
        m.pose.position.z = 0.6;
        if (m.pose.position.x == image_geometry::StereoCameraModel::MISSING_Z || m.pose.position.x < 0.1) continue;
        markers->markers.push_back(m);
      }
      marker_pub_.publish(markers);
      //cv::imshow("image", gray_big);
      //cv::waitKey(10);
    }

  }
}

void ConverterCloud::initTF()
{
  tf::TransformListener listener;
  tf::StampedTransform transform;
  try
  {
    listener.waitForTransform("/rover/base_link", "/rover/left_camera", ros::Time::now(), ros::Duration(3.0));
    listener.lookupTransform("/rover/base_link", "/rover/left_camera", ros::Time(0), transform);
  }
  catch (tf::TransformException &e)
  {
    ROS_WARN("tf not published yet. retrying...");
    return;
  }
  tf::Vector3    position = transform.getOrigin();
  tf::Quaternion quaternion = transform.getRotation();
  tf::Matrix3x3  orientation(quaternion);

  bTc_.translation() = Eigen::Vector3f(position.x(), position.y(), position.z());
  Eigen::Matrix3f R;
  for (int i = 0; i < 3; ++i)
    for (int j = 0; j < 3; ++j)
      R(i, j) = orientation[i][j];
  bTc_.linear() = R;
}

void ConverterCloud::cloud2map(const CloudPtr &cloud, nav_msgs::OccupancyGrid &map)
{
  // tmp fix-size map
  if (!map.data.size())
  {
    map.info.map_load_time = ros::Time::now();  // tmp
    map.info.resolution = 0.2;  // m/cell
    map.info.width  = 100;
    map.info.height = 100;
    map.info.origin.position.x = 0;
    map.info.origin.position.y = -map.info.resolution * map.info.height/2;
    map.info.origin.position.z = 0;
    map.info.origin.orientation.w = 1;
    map.info.origin.orientation.x = 0;
    map.info.origin.orientation.y = 0;
    map.info.origin.orientation.z = 0;
    map.data.resize(map.info.width * map.info.height);
  }

  CloudPtr new_cloud(new Cloud);
  //initTF();
  pcl::transformPointCloud(*cloud, *new_cloud, bTc_);

  /* publish */
  {
    CloudPtr cloud_disp(new Cloud);
    for (size_t i = 0; i < new_cloud->points.size(); ++i)
    {
      if (new_cloud->points[i].z < 1.5 && new_cloud->points[i].x < 5.0) 
      {
        if (new_cloud->points[i].z < -0.1) new_cloud->points[i].z = -0.14;
        cloud_disp->points.push_back(new_cloud->points[i]);
      }
    }
    cloud_disp->width = cloud_disp->points.size();
    cloud_disp->height = 1;

    sensor_msgs::PointCloud2Ptr cloud_msg(new sensor_msgs::PointCloud2);
    pcl::toROSMsg(*cloud_disp, *cloud_msg);
    cloud_msg->header.frame_id = "/rover/base_link";
    cloud_msg->header.stamp = ros::Time::now();  // tmp: refleshing time stamp
    cloud_pub_.publish(cloud_msg);
  }

  std::vector<float> zlist(new_cloud->points.size());
  for (size_t i = 0; i < new_cloud->points.size(); ++i)
  {
    zlist[i] = new_cloud->points[i].z;
    if (zlist[i] > 0.3) zlist[i] = 0.0;  //tmp
    new_cloud->points[i].z = 0.0;
  }

  float xmin =  1.0;
  float xmax =  4.0; //map.info.resolution * map.info.width;
  float ymin = -2.0; //- map.info.resolution * map.info.height / 2;
  float ymax =  2.0; // map.info.resolution * map.info.height / 2;

  OctreeSearch octree(map.info.resolution);
  octree.setInputCloud(new_cloud);
  octree.defineBoundingBox(xmin, ymin, -1, xmax, ymax, +1);
  octree.addPointsFromInputCloud();


  /* preparing markers */
  visualization_msgs::MarkerArrayPtr markers(new visualization_msgs::MarkerArray);
  visualization_msgs::Marker m;
  m.header.frame_id = "/map";
  m.header.stamp = ros::Time::now();
  m.type = visualization_msgs::Marker::CUBE;
  m.action = visualization_msgs::Marker::ADD;
  m.scale.x = 0.8 * map.info.resolution;
  m.scale.y = 0.8 * map.info.resolution;
  m.color.a = 1.0;
  m.color.r = 0.0;
  m.color.g = 1.0;
  m.color.b = 0.0;

  PointType cpt;
  cpt.z = 0.0;
  for (cpt.y = ymin; cpt.y <= ymax; cpt.y += map.info.resolution)
    for (cpt.x = xmin; cpt.x <= xmax; cpt.x += map.info.resolution)
    {
      Indices idxVec;
      float val = 0;
      if (octree.voxelSearch(cpt, idxVec) && idxVec.size() > 3)
      {
        std::vector<float> zlist_sub(idxVec.size());
        for (size_t idx = 0; idx < idxVec.size(); ++idx) zlist_sub[idx] = zlist[idxVec[idx]];
        std::sort(zlist_sub.begin(), zlist_sub.end());
        //for (size_t idx = 0; idx < idxVec.size(); ++idx) zmean += zlist[idxVec[idx]];
        //if (idxVec.size()) zmean /= idxVec.size();
        
        val = zlist_sub[3 * zlist_sub.size() / 4];
        val = std::max(val, -0.1f);  // not go under the ground!!
        m.color.a = 0.5;
        //ROS_INFO_STREAM("Neiboughers around " << cpt << " is " << idxVec.size() << " : " << val);
      }
      else
      {
        m.color.a = 1.0;
        val = 0.0;
      }
      ++m.id;
      m.pose.position.x = cpt.x;
      m.pose.position.y = cpt.y;
      m.pose.position.z = val/2;
      m.color.r = std::min(val/0.20, 1.0);
      m.color.g = std::max(1.0-val/0.30, 0.0);
      m.scale.z = val;

      markers->markers.push_back(m);

    }

  //marker_pub_.publish(markers);




  int step = map.info.width;
  for (int i = 0; i < map.info.height; ++i)
    for (int j = 0; j < map.info.width; ++j)
    {
      map.data[i * step + j] = j;
    }
}

}  // namespace stereo

// register the nodelet
#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(stereo::ConverterCloud, nodelet::Nodelet) 
